#pragma once

#include <iostream>
#include <string>
#include <sstream>
#include <shared_mutex>
#include <condition_variable>
#include <WinSock2.h>

#pragma comment(lib, "ws2_32.lib")

std::string GetErrorMsg(int code) {
    std::stringstream ss;
    if (code == WSAECONNRESET) {
        ss << "Remote Disconnected: " << code;
    } else {
        ss << "Error: " << code;
    }
    return ss.str();
}

class Err {
public:
    static SOCKET InvalidSocket;
    static int SocketError;
};

SOCKET Err::InvalidSocket = INVALID_SOCKET;
int Err::SocketError = SOCKET_ERROR;

class IpFamily {
public:
    static int v4;
    static int v6;
};

int IpFamily::v4 = AF_INET;
int IpFamily::v6 = AF_INET6;

struct _Socket {
    char *readBuffer;
    long needReadLen;
    long readIndex = 0;
    int readErr = 0;

    char *sendBuffer;
    long needWriteLen;
    long writeIndex = 0;
    bool writable = true;
    int writeErr = 0;

    WSABUF readData;
    WSABUF sendData;
    SOCKET s;

    int Read(char *data, int len) {
        readBuffer = data;
        needReadLen = len;
        readIndex = 0;

        DWORD readBytes;
        DWORD flag = 0;
        while (1) {
            readData.buf = readBuffer + readIndex;
            readData.len = needReadLen;
            if (WSARecv(s, &readData, 1, &readBytes, &flag, nullptr, nullptr) == Err::SocketError) {
                readErr = WSAGetLastError();
                return readErr;
            }
            if (readBytes == needReadLen) {
                needReadLen = 0;
                break;
            } else {
                readIndex += readBytes;
                needReadLen -= readBytes;
            }
        }
        readBuffer[len] = '\0';
        return 0;
    }

    int Write(const char *data, int len) {
        sendBuffer = new char[len];
        std::copy(data, data + len, sendBuffer);
        needWriteLen = len;
        writeIndex = 0;

        DWORD sendBytes;
        DWORD flag = 0;
        while (1) {
            sendData.buf = sendBuffer + writeIndex;
            sendData.len = needWriteLen;
            if (WSASend(s, &sendData, 1, &sendBytes, flag, nullptr, nullptr) == Err::SocketError) {
                writeErr = WSAGetLastError();
                return writeErr;
            }
            if (needWriteLen == sendBytes) {
                needWriteLen = 0;
                break;
            } else {
                writeIndex += sendBytes;
                needWriteLen -= sendBytes;
            }
        }
        delete[] sendBuffer;
        return 0;
    }

};

void Close(_Socket *so) {
    closesocket(so->s);
    delete so;
}

bool SetNonBlock(SOCKET s, u_long *isBlock) {
    if (ioctlsocket(s, FIONBIO, isBlock) == Err::SocketError) {
        return false;
    }
    return true;
}
